from prometheus_client import Counter, Histogram, Gauge
from typing import Tuple, Dict


http_requests_total = Counter(
    "http_requests_total", "HTTP requests", ["path", "method", "status"]
)

http_request_latency_ms = Histogram(
    "http_request_latency_ms", "HTTP latency", ["path", "method"]
)

e2e_latency_ms = Histogram(
    "e2e_latency_ms", "End-to-end graph invocation latency", ["path"]
)

service_available = Gauge("service_available", "Service availability")

provider_requests_total = Counter(
    "provider_requests_total", "Provider requests", ["name", "status"]
)

provider_request_latency_ms = Histogram(
    "provider_request_latency_ms", "Provider latency", ["name"]
)

provider_errors_total = Counter(
    "provider_errors_total", "Provider errors", ["name", "type"]
)

dependency_check_latency_ms = Histogram(
    "dependency_check_latency_ms", "Dependency check latency", ["name"]
)

dependency_up = Gauge(
    "dependency_up", "Dependency availability", ["name"]
)

db_query_latency_ms = Histogram(
    "db_query_latency_ms", "DB query latency", ["name"]
)

db_errors_total = Counter(
    "db_errors_total", "DB errors", ["name", "type"]
)

db_slow_queries_total = Counter(
    "db_slow_queries_total", "DB slow queries", ["name"]
)

db_queries_total = Counter(
    "db_queries_total", "DB total queries", ["name"]
)

vector_op_latency_ms = Histogram(
    "vector_op_latency_ms", "Vector DB op latency", ["name", "op"]
)

vector_errors_total = Counter(
    "vector_errors_total", "Vector DB errors", ["name", "op", "type"]
)

graph_op_latency_ms = Histogram(
    "graph_op_latency_ms", "Graph DB op latency", ["name", "op"]
)

graph_errors_total = Counter(
    "graph_errors_total", "Graph DB errors", ["name", "op", "type"]
)

nav_path_latency_ms = Histogram(
    "nav_path_latency_ms", "Navigation path computation latency", ["mode"]
)
nav_success_total = Counter(
    "nav_success_total", "Navigation success", ["mode"]
)
nav_fail_total = Counter(
    "nav_fail_total", "Navigation failures", ["mode", "reason"]
)

provider_queue_length = Gauge(
    "provider_queue_length", "Provider queue length", ["name"]
)

provider_kv_cache_hit_rate = Gauge(
    "provider_kv_cache_hit_rate", "Provider KV cache hit rate", ["name"]
)

provider_gpu_mem_used_mb = Gauge(
    "provider_gpu_mem_used_mb", "Provider GPU memory used (MB)", ["name"]
)

provider_request_tokens = Gauge(
    "provider_request_tokens", "Provider request tokens", ["name"]
)

provider_request_cost_usd = Gauge(
    "provider_request_cost_usd", "Provider request cost (USD)", ["name"]
)

robot_retry_processed_total = Counter(
    "robot_retry_processed_total", "Robot retry processed", ["result"]
)
robot_retry_queue_pops_total = Counter(
    "robot_retry_queue_pops_total", "Robot retry queue pops"
)
robot_retry_last_activity_ts = Gauge(
    "robot_retry_last_activity_ts", "Retry worker last activity epoch seconds"
)


def init_observability() -> Tuple[Counter, Histogram, Gauge]:
    return http_requests_total, http_request_latency_ms, service_available


def _aggregate_buckets(samples) -> Dict[float, float]:
    buckets: Dict[float, float] = {}
    for s in samples:
        name = s.name
        labels = s.labels
        value = s.value
        if name.endswith("_bucket"):
            le = float(labels.get("le", "inf")) if labels.get("le") != "+Inf" else float("inf")
            buckets[le] = buckets.get(le, 0.0) + float(value)
    return dict(sorted(buckets.items(), key=lambda x: x[0]))


def _aggregate_count_and_sum(samples) -> Tuple[float, float]:
    total_count = 0.0
    total_sum = 0.0
    for s in samples:
        name = s.name
        value = s.value
        if name.endswith("_count"):
            total_count += float(value)
        if name.endswith("_sum"):
            total_sum += float(value)
    return total_count, total_sum


def histogram_p95_ms(hist: Histogram) -> float:
    m = hist.collect()[0]
    buckets = _aggregate_buckets(m.samples)
    total_count, _ = _aggregate_count_and_sum(m.samples)
    if total_count <= 0:
        return 0.0
    target = 0.95 * total_count
    cum = 0.0
    for le, cnt in buckets.items():
        cum += cnt
        if cum >= target:
            return 1000.0 * le if le != float("inf") else 0.0
    return 0.0


def provider_error_rate(name: str) -> float:
    m = provider_requests_total.collect()[0]
    ok = 0.0
    total = 0.0
    for _, labels, value in m.samples:
        if labels.get("name") == name:
            total += float(value)
            if labels.get("status") == "200":
                ok += float(value)
    if total <= 0:
        return 0.0
    return round((total - ok) / total, 4)


def _aggregate_filtered_buckets(samples, label_name: str, label_value: str) -> Dict[float, float]:
    buckets: Dict[float, float] = {}
    for s in samples:
        name = s.name
        labels = s.labels
        value = s.value
        if name.endswith("_bucket") and labels.get(label_name) == label_value:
            le = float(labels.get("le", "inf")) if labels.get("le") != "+Inf" else float("inf")
            buckets[le] = buckets.get(le, 0.0) + float(value)
    return dict(sorted(buckets.items(), key=lambda x: x[0]))


def _aggregate_filtered_count_and_sum(samples, label_name: str, label_value: str) -> Tuple[float, float]:
    total_count = 0.0
    total_sum = 0.0
    for s in samples:
        name = s.name
        labels = s.labels
        value = s.value
        if labels.get(label_name) == label_value:
            if name.endswith("_count"):
                total_count += float(value)
            if name.endswith("_sum"):
                total_sum += float(value)
    return total_count, total_sum


def histogram_p95_ms_by_label(hist: Histogram, label_name: str, label_value: str) -> float:
    m = hist.collect()[0]
    buckets = _aggregate_filtered_buckets(m.samples, label_name, label_value)
    total_count, _ = _aggregate_filtered_count_and_sum(m.samples, label_name, label_value)
    if total_count <= 0:
        return 0.0
    target = 0.95 * total_count
    cum = 0.0
    for le, cnt in buckets.items():
        cum += cnt
        if cum >= target:
            return 1000.0 * le if le != float("inf") else 0.0
    return 0.0


middleware_idempotency_failures_total = Counter(
    "middleware_idempotency_failures_total", "Idempotency middleware failures", ["reason"]
)

# Confirmation prompt metrics
confirm_prompts_total = Counter(
    "confirm_prompts_total", "Confirmation prompts shown", ["intent"]
)
clarify_prompts_total = Counter(
    "clarify_prompts_total", "Clarify prompts shown", ["intent"]
)
clarify_suggestions_total = Counter(
    "clarify_suggestions_total", "Clarify suggestions generated", ["intent"]
)

# Execution success metrics
exec_success_total = Counter(
    "exec_success_total", "Successful executions", ["intent"]
)

# Execution failure metrics
exec_fail_total = Counter(
    "exec_fail_total", "Failed executions", ["intent"]
)

# Audit write failures
audit_write_failures_total = Counter(
    "audit_write_failures_total", "Audit write failures"
)

# Audit write attempts total
audit_write_total = Counter(
    "audit_write_total", "Audit write attempts total"
)
rate_limit_hits_total = Counter(
    "rate_limit_hits_total", "Rate limit hits", ["scope"]
)
blacklist_hits_total = Counter(
    "blacklist_hits_total", "Blacklist hits", ["type"]
)
risk_denials_total = Counter(
    "risk_denials_total", "Risk-based denials", ["intent", "reason"]
)

# Confirm clicks total
confirm_clicks_total = Counter(
    "confirm_clicks_total", "Confirmation clicks", ["intent"]
)

ab_exposure_total = Counter(
    "ab_exposure_total", "AB experiment exposures", ["experiment_id", "bucket"]
)

# AB-scoped execution success/failure
ab_exec_success_total = Counter(
    "ab_exec_success_total",
    "AB-scoped successful executions",
    ["experiment_id", "bucket", "intent"],
)
ab_exec_fail_total = Counter(
    "ab_exec_fail_total",
    "AB-scoped failed executions",
    ["experiment_id", "bucket", "intent"],
)

# AB-scoped clarify (confirmation) metrics
ab_clarify_prompts_total = Counter(
    "ab_clarify_prompts_total",
    "AB-scoped clarify prompts shown",
    ["experiment_id", "bucket", "intent"],
)
ab_clarify_success_total = Counter(
    "ab_clarify_success_total",
    "AB-scoped clarify successes",
    ["experiment_id", "bucket", "intent"],
)

# AB-scoped path success (by business path)
ab_path_success_total = Counter(
    "ab_path_success_total",
    "AB-scoped path successes",
    ["experiment_id", "bucket", "path"],
)

# AB assignment source distribution
ab_assignment_source_total = Counter(
    "ab_assignment_source_total",
    "AB assignment source count",
    ["experiment_id", "bucket", "source"],
)


def counter_sum_by_label(counter: Counter, label_name: str, label_value: str | None = None) -> float:
    total = 0.0
    m = counter.collect()[0]
    for name, labels, value in m.samples:
        if name.endswith("_total"):
            if label_value is None or labels.get(label_name) == label_value:
                total += float(value)
    return total


def gauge_value_by_label(gauge: Gauge, label_name: str, label_value: str) -> float:
    m = gauge.collect()[0]
    val = 0.0
    for s in m.samples:
        if s.name.endswith("_value") and s.labels.get(label_name) == label_value:
            val = float(s.value)
    return val